#include "core/packet_dispatcher.h"
#include "utils/logger.h"

#include <algorithm>
#include <chrono>

PacketDispatcher::PacketDispatcher(size_t queue_count, DispatchStrategy strategy) : m_strategy(strategy) {
    m_queues.reserve(queue_count);
    for (size_t i = 0; i < queue_count; ++i) {
        m_queues.push_back(std::make_shared<PacketQueue>());
    }

    m_stats.queue_dispatched.resize(queue_count, 0);
    m_stats.queue_sizes.resize(queue_count, 0);

    LOG_INFO("PacketDispatcher created with " << queue_count << " queues, strategy: " << static_cast<int>(strategy));
}

PacketDispatcher::~PacketDispatcher() {
    LOG_DEBUG("PacketDispatcher destroyed");
}

bool PacketDispatcher::dispatch(Packet&& packet) {
    auto start_time = std::chrono::high_resolution_clock::now();

    /// Choose the queue according to the strategy
    size_t selected_queue = 0;
    DispatchStrategy current_strategy = m_strategy.load(std::memory_order_acquire);
    switch (current_strategy) {
    case DispatchStrategy::ROUND_ROBIN:
        selected_queue = selectQueueRoundRobin();
        break;
    case DispatchStrategy::FLOW_HASH:
        selected_queue = selectQueueFlowHash(packet);
        break;
    case DispatchStrategy::LOAD_AWARE:
        selected_queue = selectQueueLoadAware();
        break;
    case DispatchStrategy::ADAPTIVE:
        selected_queue = selectQueueAdaptive(packet);
        break;
    default:
        selected_queue = selectQueueRoundRobin();
        break;
    }

    /// Attempt to push to the preferred queue
    if (m_queues[selected_queue]->try_push(std::move(packet))) {
        /// Update statistics
        {
            std::lock_guard lock(m_statsMutex);
            m_stats.total_dispatched++;
            m_stats.queue_dispatched[selected_queue]++;
        }

        // m_queues[selected_queue]->dispatched_count++;
        // updateQueueStats(selected_queue);

        /// Calculate the elapsed time
        auto end_time = std::chrono::high_resolution_clock::now();
        auto decision_time = std::chrono::duration_cast<std::chrono::nanoseconds>(end_time - start_time).count();
        m_totalDecisionTime += decision_time;

        return true;
    }

    /// If the preferred queue is full, fall back to alternative queues
    for (size_t i = 1; i < m_queues.size(); ++i) {
        size_t alt_queue = (selected_queue + i) % m_queues.size();
        if (m_queues[alt_queue]->try_push(std::move(packet))) {
            /// Update statistics
            {
                std::lock_guard lock(m_statsMutex);
                m_stats.total_dispatched++;
                m_stats.queue_dispatched[alt_queue]++;
            }

            // m_queues[alt_queue]->dispatched_count++;
            // updateQueueStats(alt_queue);

            auto end_time = std::chrono::high_resolution_clock::now();
            auto decision_time = std::chrono::duration_cast<std::chrono::nanoseconds>(end_time - start_time).count();
            m_totalDecisionTime += decision_time;

            LOG_DEBUG("Packet dispatched to alternative queue " << alt_queue);
            return true;
        }
    }

    /// All queues are full
    {
        std::lock_guard lock(m_statsMutex);
        m_stats.total_failures++;
    }

    LOG_DEBUG("All queues full, packet dropped");
    return false;
}

size_t PacketDispatcher::dispatchBulk(Packet* packets, size_t count) {
    size_t dispatched = 0;

    for (size_t i = 0; i < count; ++i) {
        if (dispatch(std::move(packets[i]))) {
            dispatched++;
        } else {
            break;
        }
    }
    return dispatched;
}

void PacketDispatcher::setStrategy(DispatchStrategy strategy) {
    auto old_strategy = m_strategy.exchange(strategy, std::memory_order_release);

    {
        std::lock_guard lock(m_statsMutex);
        m_stats.strategy_changes++;
    }

    LOG_INFO("Dispatch strategy changed from " << static_cast<int>(old_strategy)
             << " to " << static_cast<int>(strategy));
}

PacketDispatcher::PacketQueue& PacketDispatcher::getQueue(size_t index) {
    if (index >= m_queues.size()) {
        throw std::out_of_range("Queue index out of range");
    }
    return *(m_queues[index]);
}

PacketDispatcher::DispatchStats PacketDispatcher::getStats() const {
    std::lock_guard lock(m_statsMutex);
    DispatchStats stats = m_stats;

    /// Updata queue size
    for (size_t i = 0; i < m_queues.size(); ++i) {
        stats.queue_sizes[i] = m_queues[i]->size();
    }

    /// Calculate the average elapsed time
    if (stats.total_dispatched > 0) {
        stats.average_decision_time_ns = static_cast<double>(m_totalDecisionTime.load()) / stats.total_dispatched;
    }
    return stats;
}

void PacketDispatcher::resetStats() {
    std::lock_guard lock(m_statsMutex);
    m_stats = DispatchStats{};
    m_stats.queue_dispatched.resize(m_queues.size(), 0);
    m_stats.queue_sizes.resize(m_queues.size(), 0);
    m_totalDecisionTime = 0;

    LOG_DEBUG("PacketDispatcher statistics reset");
}

/// Dispatch strategy implementation
size_t PacketDispatcher::selectQueueRoundRobin() {
    return m_roundRobinIndex++ % m_queues.size();
}

size_t PacketDispatcher::selectQueueFlowHash(const Packet& packet) {
    uint64_t hash = calculateFlowHash(packet);
    return hash % m_queues.size();
}

size_t PacketDispatcher::selectQueueLoadAware() {
    return findLeastLoadedQueue();
}

size_t PacketDispatcher::selectQueueAdaptive(const Packet& packet) {
    /// Simple adaptive strategy: select based on system load
    size_t max_size = 0;
    for (const auto& queue_info : m_queues) {
        size_t size = queue_info->size();
        if (size > max_size) {
            max_size = size;
        }
    }

    /// If the maximum queue size exceeds 80% of the threshold, use load-aware scheduling
    /// 4096 is the maximum queue capacity, can be adjusted bashed on actual condition
    const size_t threshold = 4096 * 0.8;
    if (max_size > threshold) {
        return selectQueueLoadAware();
    }

    /// Otherwise, use flow hashing to preserve ordering
    return selectQueueFlowHash(packet);
}

uint32_t PacketDispatcher::calculateFlowHash(const Packet& packet) const {
    /// Simple flow hashing implementation
    /// Only support IPv4 currently
    uint32_t hash = 0;
    if (packet.isIPv4()) {
        auto ipv4 = packet.ipv4Header();
        hash ^= std::hash<uint32_t>{}(ipv4->srcAddr);
        hash ^= std::hash<uint32_t>{}(ipv4->dstAddr);
        hash ^= std::hash<uint32_t>{}(ipv4->protocol) << 16;

        if (packet.isTCP()) {
            auto tcp = packet.tcpHeader();
            hash ^= std::hash<uint32_t>{}(tcp->srcPort) << 8;
            hash ^= std::hash<uint32_t>{}(tcp->dstPort);
        } else if (packet.isUDP()) {
            auto udp = packet.udpHeader();
            hash ^= std::hash<uint32_t>{}(udp->srcPort) << 8;
            hash ^= std::hash<uint32_t>{}(udp->dstPort);
        }
    } else {
        /// Fall back to default hash
        return std::hash<const void*>{}(packet.data()) % m_queues.size();
    }

    return hash;
}

size_t PacketDispatcher::findLeastLoadedQueue() const {
    size_t best_queue = 0;
    size_t min_size = m_queues[0]->size();
    for (size_t i = 1; i < m_queues.size(); ++i) {
        size_t current_size = m_queues[i]->size();
        if (current_size < min_size) {
            min_size = current_size;
            best_queue = i;
        }
    }
    return best_queue;
}